use std::io::{self, Read, Write};
use std::net::{TcpListener, TcpStream};
use std::sync::mpsc::{self, Sender};
use std::sync::{Arc, Mutex};
use std::thread;

// 处理客户端连接
fn handle_client(mut stream: TcpStream, tx: Sender<String>) {
    let mut buffer = [0; 512];
    loop {
        match stream.read(&mut buffer) {
            Ok(n) if n == 0 => {
                println!("客户端断开连接");
                break;
            }
            Ok(n) => {
                if let Ok(message) = String::from_utf8(buffer[..n].to_vec()) {
                    println!("收到消息: {}", message.trim());
                    // 将消息广播给其他客户端
                    tx.send(message).unwrap();
                }
            }
            Err(e) => {
                eprintln!("读取错误: {}", e);
                break;
            }
        }
    }
}

// 启动聊天服务器
fn run_server() -> io::Result<()> {
    let listener = TcpListener::bind("127.0.0.1:8080")?;
    println!("服务器启动在 127.0.0.1:8080");

    let (tx, rx) = mpsc::channel::<String>();
    let client_streams = Arc::new(Mutex::new(Vec::<TcpStream>::new()));
    let client_streams_clone = Arc::clone(&client_streams);

    // 广播消息的线程
    thread::spawn(move || {
        for message in rx {
            let mut streams = client_streams_clone.lock().unwrap();
            let mut disconnected_clients = Vec::new();
            for (i, stream) in streams.iter_mut().enumerate() {
                if let Err(_) = stream.write(message.as_bytes()) {
                    disconnected_clients.push(i);
                }
            }
            // 移除断开连接的客户端
            for i in disconnected_clients.iter().rev() {
                streams.remove(*i);
            }
        }
    });

    // 接受新的客户端连接
    for stream in listener.incoming() {
        match stream {
            Ok(stream) => {
                println!("新客户端连接: {}", stream.peer_addr()?);
                let tx_clone = tx.clone();
                client_streams.lock().unwrap().push(stream.try_clone()?);
                thread::spawn(move || {
                    handle_client(stream, tx_clone);
                });
            }
            Err(e) => {
                eprintln!("接受连接错误: {}", e);
            }
        }
    }
    Ok(())
}

// 启动聊天客户端
fn run_client() -> io::Result<()> {
    let mut stream = TcpStream::connect("127.0.0.1:8080")?;
    let mut stream_clone = stream.try_clone()?;

    // 接收消息的线程
    thread::spawn(move || {
        let mut buffer = [0; 512];
        loop {
            match stream_clone.read(&mut buffer) {
                Ok(n) if n == 0 => {
                    println!("服务器已关闭");
                    break;
                }
                Ok(n) => {
                    if let Ok(message) = String::from_utf8(buffer[..n].to_vec()) {
                        println!("收到: {}", message.trim());
                    }
                }
                Err(e) => {
                    eprintln!("读取错误: {}", e);
                    break;
                }
            }
        }
    });

    // 发送消息
    let mut input = String::new();
    loop {
        input.clear();
        io::stdin().read_line(&mut input)?;
        if input.trim() == "quit" {
            break;
        }
        stream.write(input.as_bytes())?;
    }

    Ok(())
}

fn main() {
    println!("=== Rust聊天程序示例 ===\n");
    println!("请选择运行模式：");
    println!("1. 服务器模式");
    println!("2. 客户端模式");

    let mut choice = String::new();
    io::stdin().read_line(&mut choice).unwrap();

    match choice.trim() {
        "1" => {
            println!("启动服务器...");
            if let Err(e) = run_server() {
                eprintln!("服务器错误: {}", e);
            }
        }
        "2" => {
            println!("启动客户端...");
            if let Err(e) = run_client() {
                eprintln!("客户端错误: {}", e);
            }
        }
        _ => println!("无效的选择！")
    }
}